'''
Created on Jan 11, 2010

This Module is responsible for handling the firewall configuration
using the iptables tool.
'''

from os import popen
from subprocess import call

class Firewall(object) :
    '''handle packet filtering through itpables rules'''

    def listRules(self, tableName, chainName = ""):
        '''
        listRules(tableName, chainName)
            returns the rules in the selected table/chain.
        '''

        outputFile = popen("iptables -t " + tableName + " -nL " + chainName)
        rulesList = []
        for line in outputFile :
            if line.startswith('Chain') or line.startswith('target') :
                continue

            else :
                ruleOptions = line.strip().split()
                rulesList.append(ruleOptions)

        return rulesList

#####################################################################
#                                                                   #
#                    RULES-RELATED FUNCTIONS                        #
#                                                                   #
#####################################################################

    def appendRule(self, tableName, chainName, conditions, target = None):
        '''
        appendRule(tableName, chainName, conditions, target)
            Append a rule to a chain.
        '''
        if target :
            target = '-j ' + target
        else :
            target = ''

        if self.editRules(tableName, chainName, " -A ", "", conditions, target):
            return True

        else :
            return False

    def deleteRule(self, tableName, chainName, ruleID):
        '''
        deleteRule(tableName, chainName, ruleID)
            Delete a rule from a chain.
        '''
        if self.editRules(tableName, chainName, " -D ", ruleID, "", "") :
            return True

        else :
            return False

    def insertRule(self, tableName, chainName, ruleID, conditions):
        '''
        insertRule(tableName, chainName, ruleID, conditions, target)
            Insert a rule in a specific place (ruleID) in a chain.
        '''
        if call('iptables -t ' + tableName + " -I " + chainName\
                 + ' ' + ruleID + ' ' + conditions, shell = True) :
            return False

        else :
            return True

    def replaceRule(self, tableName, chainName, ruleID, conditions, target):
        '''
        replaceRule(tableName, chainName, ruleID, conditions, target)
            Replace a rule in a chain with a new rule. 
        '''
        if self.editRules(tableName, chainName, " -R ", ruleID, conditions, " -j " + target) :
            return True

        else :
            return False

    def moveRule(self, tableName, chainName, ruleID, newID, newTarget = None):
        currentRules = popen('iptables-save')
        oldRule = ''
        startRecording = False
        rules = []

        for line in currentRules :
            if startRecording == False and line.strip() == '*' + tableName :
                startRecording = True
                continue

            elif startRecording == True :
                if line.strip().startswith('-A ' + chainName) :
                    ruleTokens = line.strip().split()
                    del(ruleTokens[:2])
                    oldRule = ' '.join(ruleTokens)

                    if newTarget :
                        if ruleTokens.count('-j') == 1 and newTarget != 'noTarget' :
                            jIndex = ruleTokens.index('-j')
                            del(ruleTokens[jIndex + 1:])
                            ruleTokens.append(newTarget)

                        elif ruleTokens.count('-j') == 0 and newTarget != 'noTarget' :
                            ruleTokens.append('-j ' + newTarget)

                        elif ruleTokens.count('-j') == 1 and newTarget == 'noTarget' :
                            jIndex = ruleTokens.index('-j')
                            del(ruleTokens[jIndex:])

                    finalRule = ' '.join(ruleTokens)
                    rules.append(finalRule)

                elif line.strip() == 'COMMIT' :
                    break

        currentRules.close()
        currentRule = rules[int(ruleID) - 1]

        if not self.deleteRule(tableName, chainName, ruleID) :
            return False

        else :

            if self.insertRule(tableName, chainName, newID, currentRule) :
                return True
            else :
                self.insertRule(tableName, chainName, ruleID, oldRule)
                return False



    def editRules(self, tableName, chainName, command, commandParameters, conditions, target):

        '''
        editRules(tableName, chainName, command, commandParameters, conditions, target)
            
            Edit a rules in the selected chain.
            
            This function shouldn't be used directly by the user, instead, use other
            rules' functions and it will call this function in an appropriate way.
        '''

        if call("iptables -t " + tableName + command + chainName + " " + \
                commandParameters + " " + conditions + ' ' + target + \
                 " > /dev/null 2>&1", shell = True) :

            return False

        else :
            return True

#####################################################################
#                                                                   #
#                    CHAIN-RELATED FUNCTIONS                        #
#                                                                   #
#####################################################################
    def listChains(self, tableName):
        '''
        listChains(self, tableName):
            return list of all chain names on specific table.
        '''
        chainsListObject = popen("iptables -t " + tableName + " -nL | grep Chain")

        chainsList = []
        for line in chainsListObject :
            tempLine = line.split('(')
            chainName = tempLine[0].strip().split()[-1]

            if tempLine[1].split()[0] == 'policy' :
                policy = tempLine[1].split()[-1].split(')')[0]

            else :
                policy = 'User defined'

            chainsList.append((chainName, policy))

        return chainsList


    def flushChain(self, tableName, chainName = ""):
        '''
        flushChain(tableName, chainName = "")
            Delete all rules in the selected table/chain.
        '''

        if call("iptables -t " + tableName + " -F " + chainName , shell = True):
            return False

        else :
            return True

    def zeroChain(self, tableName, chainName = ""):
        '''
        zeroChain(tableName, chainName = "")
            Zero out all counters in the selected table/chain.
        '''

        if call("iptables -t " + tableName + " -Z " + chainName , shell = True):
            return False

        else :
            return True

    def addChain(self, tableName, chainName):
        '''
        addChain(tableName, chainName)
            Add a new user-defined chain.
        '''

        if call("iptables -t " + tableName + " -N " + chainName , shell = True):
            return False

        else :
            return True

    def delChain(self, tableName, chainName = ""):
        '''
        delChain(tableName, chainName)
            delete a user-defined chain.
        '''

        if call("iptables -t " + tableName + " -X " + chainName , shell = True):
            return False

        else :
            return True

    def renameChain(self, tableName, chainName, newName):
        '''
        renameChain(tableName, chainName, newName)
            Rename a user-defined chain.
        '''
        if call("iptables -t " + tableName + " -E " + chainName + " " + newName , shell = True):
            return False

        else :
            return True

    def setPolicy(self, tableName, chainName, target):
        '''
        setPolicy(tableName, chainName)
            Set the default action (Policy) for the selected chain.
            Note: Only built-in chains can have a policy.
        '''

        if call("iptables -t " + tableName + " -P " + chainName + " " + target , shell = True):
            return False

        else :
            return True

def main():
    # THE FOLLOWING CODE IS JUST FOR TESTING
    # CODER DISCRITION IS ADVISED :D

    x = Firewall()

    result = x.moveRule('mangle', 'OUTPUT', '6', '3')
    if result :
        print 'moved successfully'
    else :
        print 'returned error'

    print x.listRules("filter")

    if x.appendRule("filter", "INPUT", "-s 10.0.0.0/24", "ACCEPT"):
        print "rule appended successfully"
        print x.listRules("filter", "INPUT")
    else :
        print "failed to append a rule"


    if x.insertRule("filter", "INPUT", "1", "-d 172.16.0.0/16", "DROP"):
        print "rule inserted successfully"
        print x.listRules("filter", "INPUT")


    if x.deleteRule("filter", "INPUT", "1"):
        print "rule deleted successfully"
        print x.listRules("filter", "INPUT")

    else :
        print "failed to delete the rule"

    if x.replaceRule("filter", "INPUT", "1", "-d 192.168.7.0/24", "REJECT"):
        print "rule replaced successfully"
        print x.listRules("filter", "INPUT")

    else :
        print "failed to replace the rule"

    if x.flushChain("filter", "INPUT"):
        print "chain flushed successfully"
        print x.listRules("filter", "INPUT")

    else :
        print "failed to flush the chain"

    if x.zeroChain("filter", "INPUT"):
        print "chain zeroed successfully"
        print x.listRules("filter", "INPUT")

    else :
        print "failed to zero the chain"

    if x.addChain("filter", "TEST"):
        print "chain created successfully"
        print x.listRules("filter")

    else :
        print "failed to create the chain"

    if x.renameChain("filter", "TEST", "TEST2"):
        print "chain renamed successfully"
        print x.listRules("filter")

    else :
        print "failed to rename the chain"

    if x.delChain("filter", "TEST2"):
        print "chain deleted successfully"
        print x.listRules("filter")

    else :
        print "failed to delete the chain"

    if x.setPolicy("filter", "INPUT", "ACCEPT"):
        print "chain's policy changed successfully"
        print x.listRules("filter", "INPUT")

    else :
        print "failed to change the chain's policy"

if __name__ == "__main__" :

    x = Firewall()

    result = x.moveRule('mangle', 'OUTPUT', '3', '4', 'RETURN')
    if result :
        print 'moved successfully'
    else :
        print 'returned error'
    #main()
    # call the main() function here to test the module
    pass
